'''
Created on May 21, 2012

@author: tel
'''
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pylab import setp
import stats
import sys, copy, os, pickle
import scipy.stats as ss
from plothelper import Unpickle, Repickle

def OffTicks(ax, list):
    for entry in list:
        if entry=='x':
            setp(ax.get_xticklabels(), visible=False)
        elif entry=='y':
            setp(ax.get_yticklabels(), visible=False)

def TwoPanelHisto(l, r):
    fig = plt.figure()
    ax1 = fig.add_subplot(2,1,1)
    ax1.hist(l)
    ax2 = fig.add_subplot(2,1,2,sharex=ax1)
    ax2.hist(r)
    plt.show()

def OnePanelHisto(c, repickle=True, outdir='home/tel/Dropbox/remote', outfile='1panelhisto'):
    crange = (-15, 20)
    bin_width = 1
    cbins = int((crange[1]-crange[0])/bin_width)
    
    if repickle:
        Repickle(outdir, outfile, c)
    
    gs = gridspec.GridSpec(1,1)
    gs.update(wspace=0.04, hspace=0.07)
    ax1 = plt.subplot(gs[0, 0])
    for ax, data in zip((ax1), (c)):
        ax.hist(data, bins=cbins, range=crange)
    
        
    ax1.set_ylabel('surface\ncount', multialignment = 'center', labelpad=0)
    
    ax1.set_xlim(*crange)
    
    for ax in (ax1):
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            tick.label.set_rotation(45)
        for tick in ax.yaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            
    plt.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)
    
def ThreePanelHisto(l, c, r, repickle=True, outdir='home/tel/Dropbox/remote', outfile='3panelhisto'):
    lrange = (-15, 20)
    crange = (-20, 15)
    rrange = (-20, 15)
    bin_width = 1
    lbins = int((lrange[1]-lrange[0])/bin_width)
    cbins = int((crange[1]-crange[0])/bin_width)
    rbins = int((rrange[1]-rrange[0])/bin_width)
    
    if repickle:
        Repickle(outdir, outfile, l, c, r)
    
    gs = gridspec.GridSpec(3,2)
    gs.update(wspace=0.04, hspace=0.07)
    ax1 = plt.subplot(gs[0, 0])
    ax2 = plt.subplot(gs[0, 1])
    ax3 = plt.subplot(gs[1, 0])
    
    for ax, data, bins, range in zip((ax1, ax2, ax3), (l, c, r), (lbins,cbins,rbins),(lrange,crange,rrange)):
        ax.hist(data, bins=bins, range=range)
        
    ax1.set_ylabel('surface\ncount', multialignment = 'center', labelpad=0)
    ax2.set_ylabel('mildly buried\ncount', multialignment = 'center', labelpad=0)
    ax3.set_ylabel('highly buried\ncount', multialignment = 'center', labelpad=0)
    
#    ax1.set_ylabel('Low ${\Delta G}_{dsol}$\n(surface) residue count', multialignment = 'center', labelpad=0)
#    ax3.set_ylabel('Mid ${\Delta G}_{dsol}$\nresidue count', multialignment = 'center', labelpad=5)
#    ax5.set_ylabel('High ${\Delta G}_{dsol}$\n(buried) residue count', multialignment = 'center',labelpad=10)
#    ax6s.set_ylabel('Mutant High ${\Delta G}_{dsol}$\n(buried) residue count', multialignment = 'center',labelpad=10)
#    ax5.set_xlabel('$\Delta pK_{a}$')
#    ax6.set_xlabel('$\Delta\Delta pK_{a}$')
    
    ax1.set_xlim(*lrange)
    ax2.set_xlim(*crange)
    ax3.set_xlim(*rrange)
    
    for ax in (ax1,ax2,ax3):
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            tick.label.set_rotation(45)
        for tick in ax.yaxis.get_major_ticks():
            tick.label.set_fontsize('small')

    plt.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)

def SixPanelHisto(ul, ur, ml, mr, ll, lr, lls, lrs, repickle=True, outdir='home/tel/Dropbox/remote', outfile='6panelhisto'):
    lrange = (-15, 20)
    rrange = (-20, 15)
    bin_width = 1
    lbins = int((lrange[1]-lrange[0])/bin_width)
    rbins = int((rrange[1]-rrange[0])/bin_width)
    
    if repickle:
        Repickle(outdir, outfile, ul, ur, ml, mr, ll, lr, lls, lrs)
    
    gs = gridspec.GridSpec(3,2)
    gs.update(wspace=0.04, hspace=0.07)
    ax5 = plt.subplot(gs[2, 0])
    ax6 = plt.subplot(gs[2, 1], sharey=ax5)
    ax4 = plt.subplot(gs[1, 1], sharex=ax6)
    ax3 = plt.subplot(gs[1, 0], sharex=ax5, sharey=ax4)
    ax1 = plt.subplot(gs[0, 0], sharex=ax5)
    ax2 = plt.subplot(gs[0, 1], sharex=ax6, sharey=ax1)
    
    ax5s = ax5.twinx()
    ax6s = ax6.twinx()
    
    for ax, data in zip((ax1, ax3, ax5), (ul, ml, ll)):
        ax.hist(data, bins=lbins, range=lrange, color='r')
    for ax, data in zip((ax2, ax4, ax6), (ur, mr, lr)):
        ax.hist(data, bins=rbins, range=rrange, color='r')
    
#    for ax, data in zip((ax5s, ax6s), (lls, lrs)):
#        ax.set_ylim(0, 30)
#        ax.hist(data, bins=lbins, range=lrange, color='g', alpha=.75)
        
    ax1.set_ylabel('surface\ncount', multialignment = 'center', labelpad=0)
    ax3.set_ylabel('mildly buried\ncount', multialignment = 'center', labelpad=5)
    ax5.set_ylabel('highly buried\ncount', multialignment = 'center',labelpad=10)
    ax6s.set_ylabel('highly buried\nmutant count', multialignment = 'center',labelpad=10)
    ax5.set_xlabel('$\Delta pK_{a}$')
    ax6.set_xlabel('$\Delta\Delta pK_{a}$')
    
#    ax1.set_ylabel('Low ${\Delta G}_{dsol}$\n(surface) residue count', multialignment = 'center', labelpad=0)
#    ax3.set_ylabel('Mid ${\Delta G}_{dsol}$\nresidue count', multialignment = 'center', labelpad=5)
#    ax5.set_ylabel('High ${\Delta G}_{dsol}$\n(buried) residue count', multialignment = 'center',labelpad=10)
#    ax6s.set_ylabel('Mutant High ${\Delta G}_{dsol}$\n(buried) residue count', multialignment = 'center',labelpad=10)
#    ax5.set_xlabel('$\Delta pK_{a}$')
#    ax6.set_xlabel('$\Delta\Delta pK_{a}$')
    
    ax5.set_xlim(*lrange)
    ax6.set_xlim(*rrange)
    
    for ax in (ax1,ax2,ax3,ax4,ax5,ax6,ax5s,ax6s):
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            tick.label.set_rotation(45)
        for tick in ax.yaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            
    OffTicks(ax1, ['x'])
    OffTicks(ax2, ['x','y'])
    OffTicks(ax3, ['x'])
    OffTicks(ax4, ['x','y'])
    OffTicks(ax5, [])
    OffTicks(ax6, ['y'])
    OffTicks(ax5s, ['y'])
    plt.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)

if __name__=="__main__":
    workdir = sys.argv[1]
    #Unpickle(indir=workdir, infile='dpka_ddpka_histos', outdir=workdir, outfile='dpka_ddpka_histos', argn=8, plotter=SixPanelHisto)
    #Unpickle(indir=workdir, infile='dpka_ddpka-lowdsol_middsol_highdsol-derk-histos', outdir=workdir, outfile='dpka_ddpka-lowdsol_middsol_highdsol-derk-histos', argn=8, plotter=SixPanelHisto)
    Unpickle(indir=workdir, infile='dpka_ddpka-lowdsol_middsol_highdsol-his-histos', outdir=workdir, outfile='dpka_ddpka-lowdsol_middsol_highdsol-his-histos', argn=8, plotter=SixPanelHisto)
#    data = []
#    for i in range(6):
#        data.append(np.random.rand(10,1)*(np.random.rand()*100))
#    SixPanelHisto(data[0], data[1], data[2], data[3], data[4], data[5])
